#!/usr/bin/env python
# Simple document querier
# David Coles (239715), 2010

import math

import index

SNIPPET_BEFORE = 20
SNIPPET_AFTER = 50

class BM25(object):
    """ Uses a SimpleIndex to score a document using Okapi BM 25.
    """
    def __init__(self, simple_index, k1=2.0, b=0.75):
        self.index = simple_index
        self.k1 = k1
        self.b = b

    def idf_weight(self, term):
        """ Calculates the inverse document frequency weight of a term
        """
        # To make the formula clear, we set up local variables in terms of the 
        # mathematical description
        N = self.index.N()
        n = self.index.doc_freq

        return math.log((N - n(term) + 0.5)/(n(term) + 0.5))

    def score(self, document_id, query):
        """ Calculates the Okapi BM25 score for a document
        """
        # To make the formula clear, we set up local variables in terms of the 
        # mathematical description
        f = self.index.term_freq
        D = document_id
        k1 = self.k1
        b = self.b
        l = self.index.length(document_id)
        avg_l = self.index.avg_length()

        score = 0
        for term in query:
            idf = self.idf_weight(term)
            s = idf * (f(term,D) * (k1 + 1)) / (
                    f(term,D) + k1 * (1 - b + b*l/avg_l))
            # Ignore negative scored terms (common terms)
            if s > 0.0:
                score += s

        return score

    def search(self, query):
        """ Returns a ranked list of (score, document_id) tuples for all 
        documents indexed.
        """
        # Score all documents
        scores = {}
        for document in self.index.documents():
            scores[document] = self.score(document, query)

        sorted_docs = sorted(scores, key=lambda k: scores[k], reverse=True)

        # Return a list of (score, doc) tuples sorted by score
        return [(scores[doc], doc) for doc in sorted_docs]


def snippet(text, positions):
    """ Returns a list of snippets that span all the positions in text with at 
    least SNIPPET_BEFORE characters before and SNIPPET_AFTER characters after.
    """

    if positions:
        # Must be sorted
        positions = sorted(positions)

        spans = []
        start = max(0, positions[0]-SNIPPET_BEFORE)
        prev = positions[0]

        for p in positions[1:]:
            # If span does not overlap print and start new span
            if p-SNIPPET_BEFORE >= prev+SNIPPET_AFTER:
                spans.append(text[start:prev+SNIPPET_AFTER])
                start = p-SNIPPET_BEFORE
            prev = p

        # Close the final span
        spans.append(text[start:prev+SNIPPET_AFTER])
        return spans
    else:
        # No matching offsets
        return ''


if __name__ == "__main__":
    # Hard coded query
    query = ['approbation', 'demain', 'matin']

    # Load index
    with open('french.dat', 'r') as f:
        index = index.SimpleIndex(f)

    # Apply the BM25 model to our index
    bm25 = BM25(index)

    # Score all documents
    scores = {}
    for document in index.documents():
        scores[document] = bm25.score(document, query)

    sorted_docs = sorted(scores, key=lambda k: scores[k], reverse=True)

    print "Query for: " + ' '.join(query)
    for doc in sorted_docs[:10]:
        print "%10.2f: %s"%(scores[doc], doc)

